import onnx
from onnxsim import simplify


def optimize_onnx(input_path, output_path):
    """进一步优化 ONNX 模型"""
    # 加载模型
    model = onnx.load(input_path)

    # 模型简化
    model_simp, check = simplify(model)

    # 保存优化后模型
    onnx.save(model_simp, output_path)

    # 打印优化结果
    orig_size = len(model.SerializeToString()) / 1024 ** 2
    new_size = len(model_simp.SerializeToString()) / 1024 ** 2
    print(f"📦 ONNX 优化完成: {orig_size:.2f}MB → {new_size:.2f}MB (-{100 * (1 - new_size / orig_size):.1f}%)")
    print(f"✅ 优化模型保存至: {output_path}")


# 示例使用
optimize_onnx('../models/yolov8n.onnx', '../models/yolov8n_optimized.onnx')